{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5aElYAKlV2Mi"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wmYJlt6LWVOU"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-8q8rRRWcp6"
      },
      "source": [
        "# TensorFlow Addons Networks : Sequence-to-Sequence NMT with Attention Mechanism\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/addons/tutorials/networks_seq2seq_nmt\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_nmt.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/addons/blob/master/docs/tutorials/networks_seq2seq_nmt.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "      <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/addons/docs/tutorials/networks_seq2seq_nmt.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9n0dcDw1Wszw"
      },
      "source": [
        "## Overview\n",
        "This notebook gives a brief introduction into the ***Sequence to Sequence Model Architecture***\n",
        "In this noteboook you broadly cover four essential topics necessary for Neural Machine Translation:\n",
        "\n",
        "\n",
        "* **Data cleaning**\n",
        "* **Data preparation**\n",
        "* **Neural Translation Model with Attention**\n",
        "* **Final Translation with ```tf.addons.seq2seq.BasicDecoder``` and ```tf.addons.seq2seq.BeamSearchDecoder```** \n",
        "\n",
        "The basic idea behind such a model though, is only the encoder-decoder architecture. These networks are usually used for a variety of tasks like text-summerization, Machine translation, Image Captioning, etc. This tutorial provideas a hands-on understanding of the concept, explaining the technical jargons wherever necessary. You focus on the task of Neural Machine Translation (NMT) which was the very first testbed for seq2seq models.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MpySVYWJhxaV"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_kxfdP4hJUPB"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting tensorflow-addons==0.11.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b3/f8/d6fca180c123f2851035c4493690662ebdad0849a9059d56035434bff5c9/tensorflow_addons-0.11.2-cp36-cp36m-manylinux2010_x86_64.whl (1.1MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1MB 4.4MB/s \n",
            "\u001b[?25hRequirement already satisfied: typeguard>=2.7 in /usr/local/lib/python3.6/dist-packages (from tensorflow-addons==0.11.2) (2.7.1)\n",
            "Installing collected packages: tensorflow-addons\n",
            "  Found existing installation: tensorflow-addons 0.11.0\n",
            "    Uninstalling tensorflow-addons-0.11.0:\n",
            "      Successfully uninstalled tensorflow-addons-0.11.0\n",
            "Successfully installed tensorflow-addons-0.11.2\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow-addons==0.11.2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnxXKDjq3jEL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as ticker\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "import unicodedata\n",
        "import re\n",
        "import numpy as np\n",
        "import os\n",
        "import io\n",
        "import time\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ii_vg-XNXTil"
      },
      "source": [
        "## Data Cleaning and Data Preparation \n",
        "\n",
        "You'll use a language dataset provided by http://www.manythings.org/anki/. This dataset contains language translation pairs in the format:\n",
        "\n",
        "---\n",
        "      May I borrow this book?    ¿Puedo tomar prestado este libro?\n",
        "---\n",
        "\n",
        "\n",
        "There are a variety of languages available, but you'll use the English-Spanish dataset. After downloading the dataset, here are the steps you'll take to prepare the data:\n",
        "\n",
        "\n",
        "1. Add a start and end token to each sentence.\n",
        "2. Clean the sentences by removing special characters.\n",
        "3. Create a Vocabulary with word index (mapping from word → id) and reverse word index (mapping from id → word).\n",
        "5. Pad each sentence to a maximum length. (Why? you need to fix the maximum length for the inputs to recurrent encoders)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PvRnGWnvXm6l"
      },
      "outputs": [],
      "source": [
        "def download_nmt():\n",
        "    path_to_zip = tf.keras.utils.get_file(\n",
        "    'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',\n",
        "    extract=True)\n",
        "\n",
        "    path_to_file = os.path.dirname(path_to_zip)+\"/spa-eng/spa.txt\"\n",
        "    return path_to_file\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NFKB2c_tX4wU"
      },
      "source": [
        "### Define a NMTDataset class with necessary functions to follow Step 1 to Step 4. \n",
        "The ```call()``` will return:\n",
        "1. ```train_dataset```  and ```val_dataset``` : ```tf.data.Dataset``` objects\n",
        "2. ```inp_lang_tokenizer``` and ```targ_lang_tokenizer``` : ```tf.keras.preprocessing.text.Tokenizer``` objects "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JMAHz7kJXc5N"
      },
      "outputs": [],
      "source": [
        "class NMTDataset:\n",
        "    def __init__(self, problem_type='en-spa'):\n",
        "        self.problem_type = 'en-spa'\n",
        "        self.inp_lang_tokenizer = None\n",
        "        self.targ_lang_tokenizer = None\n",
        "    \n",
        "\n",
        "    def unicode_to_ascii(self, s):\n",
        "        return ''.join(c for c in unicodedata.normalize('NFD', s) if unicodedata.category(c) != 'Mn')\n",
        "\n",
        "    ## Step 1 and Step 2 \n",
        "    def preprocess_sentence(self, w):\n",
        "        w = self.unicode_to_ascii(w.lower().strip())\n",
        "\n",
        "        # creating a space between a word and the punctuation following it\n",
        "        # eg: \"he is a boy.\" => \"he is a boy .\"\n",
        "        # Reference:- https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation\n",
        "        w = re.sub(r\"([?.!,¿])\", r\" \\1 \", w)\n",
        "        w = re.sub(r'[\" \"]+', \" \", w)\n",
        "\n",
        "        # replacing everything with space except (a-z, A-Z, \".\", \"?\", \"!\", \",\")\n",
        "        w = re.sub(r\"[^a-zA-Z?.!,¿]+\", \" \", w)\n",
        "\n",
        "        w = w.strip()\n",
        "\n",
        "        # adding a start and an end token to the sentence\n",
        "        # so that the model know when to start and stop predicting.\n",
        "        w = '<start> ' + w + ' <end>'\n",
        "        return w\n",
        "    \n",
        "    def create_dataset(self, path, num_examples):\n",
        "        # path : path to spa-eng.txt file\n",
        "        # num_examples : Limit the total number of training example for faster training (set num_examples = len(lines) to use full data)\n",
        "        lines = io.open(path, encoding='UTF-8').read().strip().split('\\n')\n",
        "        word_pairs = [[self.preprocess_sentence(w) for w in l.split('\\t')]  for l in lines[:num_examples]]\n",
        "\n",
        "        return zip(*word_pairs)\n",
        "\n",
        "    # Step 3 and Step 4\n",
        "    def tokenize(self, lang):\n",
        "        # lang = list of sentences in a language\n",
        "        \n",
        "        # print(len(lang), \"example sentence: {}\".format(lang[0]))\n",
        "        lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(filters='', oov_token='<OOV>')\n",
        "        lang_tokenizer.fit_on_texts(lang)\n",
        "\n",
        "        ## tf.keras.preprocessing.text.Tokenizer.texts_to_sequences converts string (w1, w2, w3, ......, wn) \n",
        "        ## to a list of correspoding integer ids of words (id_w1, id_w2, id_w3, ...., id_wn)\n",
        "        tensor = lang_tokenizer.texts_to_sequences(lang) \n",
        "\n",
        "        ## tf.keras.preprocessing.sequence.pad_sequences takes argument a list of integer id sequences \n",
        "        ## and pads the sequences to match the longest sequences in the given input\n",
        "        tensor = tf.keras.preprocessing.sequence.pad_sequences(tensor, padding='post')\n",
        "\n",
        "        return tensor, lang_tokenizer\n",
        "\n",
        "    def load_dataset(self, path, num_examples=None):\n",
        "        # creating cleaned input, output pairs\n",
        "        targ_lang, inp_lang = self.create_dataset(path, num_examples)\n",
        "\n",
        "        input_tensor, inp_lang_tokenizer = self.tokenize(inp_lang)\n",
        "        target_tensor, targ_lang_tokenizer = self.tokenize(targ_lang)\n",
        "\n",
        "        return input_tensor, target_tensor, inp_lang_tokenizer, targ_lang_tokenizer\n",
        "\n",
        "    def call(self, num_examples, BUFFER_SIZE, BATCH_SIZE):\n",
        "        file_path = download_nmt()\n",
        "        input_tensor, target_tensor, self.inp_lang_tokenizer, self.targ_lang_tokenizer = self.load_dataset(file_path, num_examples)\n",
        "        \n",
        "        input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.2)\n",
        "\n",
        "        train_dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train))\n",
        "        train_dataset = train_dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)\n",
        "\n",
        "        val_dataset = tf.data.Dataset.from_tensor_slices((input_tensor_val, target_tensor_val))\n",
        "        val_dataset = val_dataset.batch(BATCH_SIZE, drop_remainder=True)\n",
        "\n",
        "        return train_dataset, val_dataset, self.inp_lang_tokenizer, self.targ_lang_tokenizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EIW4NVBmJ25k"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 32000\n",
        "BATCH_SIZE = 64\n",
        "# Let's limit the #training examples for faster training\n",
        "num_examples = 30000\n",
        "\n",
        "dataset_creator = NMTDataset('en-spa')\n",
        "train_dataset, val_dataset, inp_lang, targ_lang = dataset_creator.call(num_examples, BUFFER_SIZE, BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w2lCTy4vKOkB"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(TensorShape([64, 16]), TensorShape([64, 11]))"
            ]
          },
          "execution_count": 7,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "example_input_batch, example_target_batch = next(iter(train_dataset))\n",
        "example_input_batch.shape, example_target_batch.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgCLkfv5uO3d"
      },
      "source": [
        "### Some important parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqHsArVZ3jFS"
      },
      "outputs": [],
      "source": [
        "vocab_inp_size = len(inp_lang.word_index)+1\n",
        "vocab_tar_size = len(targ_lang.word_index)+1\n",
        "max_length_input = example_input_batch.shape[1]\n",
        "max_length_output = example_target_batch.shape[1]\n",
        "\n",
        "embedding_dim = 256\n",
        "units = 1024\n",
        "steps_per_epoch = num_examples//BATCH_SIZE\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g-yY9c6aIu1h"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "max_length_spanish, max_length_english, vocab_size_spanish, vocab_size_english\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "(16, 11, 9415, 4936)"
            ]
          },
          "execution_count": 9,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "print(\"max_length_spanish, max_length_english, vocab_size_spanish, vocab_size_english\")\n",
        "max_length_input, max_length_output, vocab_inp_size, vocab_tar_size"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nZ2rI24i3jFg"
      },
      "outputs": [],
      "source": [
        "##### \n",
        "\n",
        "class Encoder(tf.keras.Model):\n",
        "  def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):\n",
        "    super(Encoder, self).__init__()\n",
        "    self.batch_sz = batch_sz\n",
        "    self.enc_units = enc_units\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
        "\n",
        "    ##-------- LSTM layer in Encoder ------- ##\n",
        "    self.lstm_layer = tf.keras.layers.LSTM(self.enc_units,\n",
        "                                   return_sequences=True,\n",
        "                                   return_state=True,\n",
        "                                   recurrent_initializer='glorot_uniform')\n",
        "    \n",
        "\n",
        "\n",
        "  def call(self, x, hidden):\n",
        "    x = self.embedding(x)\n",
        "    output, h, c = self.lstm_layer(x, initial_state = hidden)\n",
        "    return output, h, c\n",
        "\n",
        "  def initialize_hidden_state(self):\n",
        "    return [tf.zeros((self.batch_sz, self.enc_units)), tf.zeros((self.batch_sz, self.enc_units))] "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "60gSVh05Jl6l"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Encoder output shape: (batch size, sequence length, units) (64, 16, 1024)\n",
            "Encoder h vecotr shape: (batch size, units) (64, 1024)\n",
            "Encoder c vector shape: (batch size, units) (64, 1024)\n"
          ]
        }
      ],
      "source": [
        "## Test Encoder Stack\n",
        "\n",
        "encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)\n",
        "\n",
        "\n",
        "# sample input\n",
        "sample_hidden = encoder.initialize_hidden_state()\n",
        "sample_output, sample_h, sample_c = encoder(example_input_batch, sample_hidden)\n",
        "print ('Encoder output shape: (batch size, sequence length, units) {}'.format(sample_output.shape))\n",
        "print ('Encoder h vecotr shape: (batch size, units) {}'.format(sample_h.shape))\n",
        "print ('Encoder c vector shape: (batch size, units) {}'.format(sample_c.shape))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yJ_B3mhW3jFk"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.Model):\n",
        "  def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz, attention_type='luong'):\n",
        "    super(Decoder, self).__init__()\n",
        "    self.batch_sz = batch_sz\n",
        "    self.dec_units = dec_units\n",
        "    self.attention_type = attention_type\n",
        "    \n",
        "    # Embedding Layer\n",
        "    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n",
        "    \n",
        "    #Final Dense layer on which softmax will be applied\n",
        "    self.fc = tf.keras.layers.Dense(vocab_size)\n",
        "\n",
        "    # Define the fundamental cell for decoder recurrent structure\n",
        "    self.decoder_rnn_cell = tf.keras.layers.LSTMCell(self.dec_units)\n",
        "   \n",
        "\n",
        "\n",
        "    # Sampler\n",
        "    self.sampler = tfa.seq2seq.sampler.TrainingSampler()\n",
        "\n",
        "    # Create attention mechanism with memory = None\n",
        "    self.attention_mechanism = self.build_attention_mechanism(self.dec_units, \n",
        "                                                              None, self.batch_sz*[max_length_input], self.attention_type)\n",
        "\n",
        "    # Wrap attention mechanism with the fundamental rnn cell of decoder\n",
        "    self.rnn_cell = self.build_rnn_cell(batch_sz)\n",
        "\n",
        "    # Define the decoder with respect to fundamental rnn cell\n",
        "    self.decoder = tfa.seq2seq.BasicDecoder(self.rnn_cell, sampler=self.sampler, output_layer=self.fc)\n",
        "\n",
        "    \n",
        "  def build_rnn_cell(self, batch_sz):\n",
        "    rnn_cell = tfa.seq2seq.AttentionWrapper(self.decoder_rnn_cell, \n",
        "                                  self.attention_mechanism, attention_layer_size=self.dec_units)\n",
        "    return rnn_cell\n",
        "\n",
        "  def build_attention_mechanism(self, dec_units, memory, memory_sequence_length, attention_type='luong'):\n",
        "    # ------------- #\n",
        "    # typ: Which sort of attention (Bahdanau, Luong)\n",
        "    # dec_units: final dimension of attention outputs \n",
        "    # memory: encoder hidden states of shape (batch_size, max_length_input, enc_units)\n",
        "    # memory_sequence_length: 1d array of shape (batch_size) with every element set to max_length_input (for masking purpose)\n",
        "\n",
        "    if(attention_type=='bahdanau'):\n",
        "      return tfa.seq2seq.BahdanauAttention(units=dec_units, memory=memory, memory_sequence_length=memory_sequence_length)\n",
        "    else:\n",
        "      return tfa.seq2seq.LuongAttention(units=dec_units, memory=memory, memory_sequence_length=memory_sequence_length)\n",
        "\n",
        "  def build_initial_state(self, batch_sz, encoder_state, Dtype):\n",
        "    decoder_initial_state = self.rnn_cell.get_initial_state(batch_size=batch_sz, dtype=Dtype)\n",
        "    decoder_initial_state = decoder_initial_state.clone(cell_state=encoder_state)\n",
        "    return decoder_initial_state\n",
        "\n",
        "\n",
        "  def call(self, inputs, initial_state):\n",
        "    x = self.embedding(inputs)\n",
        "    outputs, _, _ = self.decoder(x, initial_state=initial_state, sequence_length=self.batch_sz*[max_length_output-1])\n",
        "    return outputs\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DaiO0Z6_Ml1c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Decoder Outputs Shape:  (64, 10, 4936)\n"
          ]
        }
      ],
      "source": [
        "# Test decoder stack\n",
        "\n",
        "decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE, 'luong')\n",
        "sample_x = tf.random.uniform((BATCH_SIZE, max_length_output))\n",
        "decoder.attention_mechanism.setup_memory(sample_output)\n",
        "initial_state = decoder.build_initial_state(BATCH_SIZE, [sample_h, sample_c], tf.float32)\n",
        "\n",
        "\n",
        "sample_decoder_outputs = decoder(sample_x, initial_state)\n",
        "\n",
        "print(\"Decoder Outputs Shape: \", sample_decoder_outputs.rnn_output.shape)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ch_71VbIRfK"
      },
      "source": [
        "## Define the optimizer and the loss function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WmTHr5iV3jFr"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.Adam()\n",
        "\n",
        "\n",
        "def loss_function(real, pred):\n",
        "  # real shape = (BATCH_SIZE, max_length_output)\n",
        "  # pred shape = (BATCH_SIZE, max_length_output, tar_vocab_size )\n",
        "  cross_entropy = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')\n",
        "  loss = cross_entropy(y_true=real, y_pred=pred)\n",
        "  mask = tf.logical_not(tf.math.equal(real,0))   #output 0 for y=0 else output 1\n",
        "  mask = tf.cast(mask, dtype=loss.dtype)  \n",
        "  loss = mask* loss\n",
        "  loss = tf.reduce_mean(loss)\n",
        "  return loss  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DMVWzzsfNl4e"
      },
      "source": [
        "## Checkpoints (Object-based saving)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zj8bXQTgNwrF"
      },
      "outputs": [],
      "source": [
        "checkpoint_dir = './training_checkpoints'\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
        "checkpoint = tf.train.Checkpoint(optimizer=optimizer,\n",
        "                                 encoder=encoder,\n",
        "                                 decoder=decoder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Bw95utNiFHa"
      },
      "source": [
        "## One train_step operations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sC9ArXSsVfqn"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(inp, targ, enc_hidden):\n",
        "  loss = 0\n",
        "\n",
        "  with tf.GradientTape() as tape:\n",
        "    enc_output, enc_h, enc_c = encoder(inp, enc_hidden)\n",
        "\n",
        "\n",
        "    dec_input = targ[ : , :-1 ] # Ignore <end> token\n",
        "    real = targ[ : , 1: ]         # ignore <start> token\n",
        "\n",
        "    # Set the AttentionMechanism object with encoder_outputs\n",
        "    decoder.attention_mechanism.setup_memory(enc_output)\n",
        "\n",
        "    # Create AttentionWrapperState as initial_state for decoder\n",
        "    decoder_initial_state = decoder.build_initial_state(BATCH_SIZE, [enc_h, enc_c], tf.float32)\n",
        "    pred = decoder(dec_input, decoder_initial_state)\n",
        "    logits = pred.rnn_output\n",
        "    loss = loss_function(real, logits)\n",
        "\n",
        "  variables = encoder.trainable_variables + decoder.trainable_variables\n",
        "  gradients = tape.gradient(loss, variables)\n",
        "  optimizer.apply_gradients(zip(gradients, variables))\n",
        "\n",
        "  return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pey8eb9piMMg"
      },
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ddefjBMa3jF0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1 Batch 0 Loss 5.1692\n",
            "Epoch 1 Batch 100 Loss 2.2288\n",
            "Epoch 1 Batch 200 Loss 1.9930\n",
            "Epoch 1 Batch 300 Loss 1.7783\n",
            "Epoch 1 Loss 1.6975\n",
            "Time taken for 1 epoch 37.26002788543701 sec\n",
            "\n",
            "Epoch 2 Batch 0 Loss 1.6408\n",
            "Epoch 2 Batch 100 Loss 1.5767\n",
            "Epoch 2 Batch 200 Loss 1.4054\n",
            "Epoch 2 Batch 300 Loss 1.3755\n",
            "Epoch 2 Loss 1.1412\n",
            "Time taken for 1 epoch 30.0094051361084 sec\n",
            "\n",
            "Epoch 3 Batch 0 Loss 1.0296\n",
            "Epoch 3 Batch 100 Loss 1.0306\n",
            "Epoch 3 Batch 200 Loss 1.0675\n",
            "Epoch 3 Batch 300 Loss 0.9574\n",
            "Epoch 3 Loss 0.8037\n",
            "Time taken for 1 epoch 28.983767986297607 sec\n",
            "\n",
            "Epoch 4 Batch 0 Loss 0.5923\n",
            "Epoch 4 Batch 100 Loss 0.7533\n",
            "Epoch 4 Batch 200 Loss 0.7397\n",
            "Epoch 4 Batch 300 Loss 0.6779\n",
            "Epoch 4 Loss 0.5419\n",
            "Time taken for 1 epoch 29.649972200393677 sec\n",
            "\n",
            "Epoch 5 Batch 0 Loss 0.4320\n",
            "Epoch 5 Batch 100 Loss 0.4349\n",
            "Epoch 5 Batch 200 Loss 0.4686\n",
            "Epoch 5 Batch 300 Loss 0.4748\n",
            "Epoch 5 Loss 0.3827\n",
            "Time taken for 1 epoch 29.06334638595581 sec\n",
            "\n",
            "Epoch 6 Batch 0 Loss 0.3422\n",
            "Epoch 6 Batch 100 Loss 0.3052\n",
            "Epoch 6 Batch 200 Loss 0.3288\n",
            "Epoch 6 Batch 300 Loss 0.3216\n",
            "Epoch 6 Loss 0.2814\n",
            "Time taken for 1 epoch 29.57170796394348 sec\n",
            "\n",
            "Epoch 7 Batch 0 Loss 0.2129\n",
            "Epoch 7 Batch 100 Loss 0.2382\n",
            "Epoch 7 Batch 200 Loss 0.2406\n",
            "Epoch 7 Batch 300 Loss 0.2792\n",
            "Epoch 7 Loss 0.2162\n",
            "Time taken for 1 epoch 28.95500087738037 sec\n",
            "\n",
            "Epoch 8 Batch 0 Loss 0.2073\n",
            "Epoch 8 Batch 100 Loss 0.2095\n",
            "Epoch 8 Batch 200 Loss 0.1962\n",
            "Epoch 8 Batch 300 Loss 0.1879\n",
            "Epoch 8 Loss 0.1794\n",
            "Time taken for 1 epoch 29.70877432823181 sec\n",
            "\n",
            "Epoch 9 Batch 0 Loss 0.1517\n",
            "Epoch 9 Batch 100 Loss 0.2231\n",
            "Epoch 9 Batch 200 Loss 0.2203\n",
            "Epoch 9 Batch 300 Loss 0.2282\n",
            "Epoch 9 Loss 0.1496\n",
            "Time taken for 1 epoch 29.20821261405945 sec\n",
            "\n",
            "Epoch 10 Batch 0 Loss 0.1204\n",
            "Epoch 10 Batch 100 Loss 0.1370\n",
            "Epoch 10 Batch 200 Loss 0.1778\n",
            "Epoch 10 Batch 300 Loss 0.2069\n",
            "Epoch 10 Loss 0.1316\n",
            "Time taken for 1 epoch 29.576894283294678 sec\n",
            "\n"
          ]
        }
      ],
      "source": [
        "EPOCHS = 10\n",
        "\n",
        "for epoch in range(EPOCHS):\n",
        "  start = time.time()\n",
        "\n",
        "  enc_hidden = encoder.initialize_hidden_state()\n",
        "  total_loss = 0\n",
        "  # print(enc_hidden[0].shape, enc_hidden[1].shape)\n",
        "\n",
        "  for (batch, (inp, targ)) in enumerate(train_dataset.take(steps_per_epoch)):\n",
        "    batch_loss = train_step(inp, targ, enc_hidden)\n",
        "    total_loss += batch_loss\n",
        "\n",
        "    if batch % 100 == 0:\n",
        "      print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,\n",
        "                                                   batch,\n",
        "                                                   batch_loss.numpy()))\n",
        "  # saving (checkpoint) the model every 2 epochs\n",
        "  if (epoch + 1) % 2 == 0:\n",
        "    checkpoint.save(file_prefix = checkpoint_prefix)\n",
        "\n",
        "  print('Epoch {} Loss {:.4f}'.format(epoch + 1,\n",
        "                                      total_loss / steps_per_epoch))\n",
        "  print('Time taken for 1 epoch {} sec\\n'.format(time.time() - start))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mU3Ce8M6I3rz"
      },
      "source": [
        "## Use tf-addons BasicDecoder for decoding\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "id": "EbQpyYs13jF_"
      },
      "outputs": [],
      "source": [
        "def evaluate_sentence(sentence):\n",
        "  sentence = dataset_creator.preprocess_sentence(sentence)\n",
        "\n",
        "  inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]\n",
        "  inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],\n",
        "                                                          maxlen=max_length_input,\n",
        "                                                          padding='post')\n",
        "  inputs = tf.convert_to_tensor(inputs)\n",
        "  inference_batch_size = inputs.shape[0]\n",
        "  result = ''\n",
        "\n",
        "  enc_start_state = [tf.zeros((inference_batch_size, units)), tf.zeros((inference_batch_size,units))]\n",
        "  enc_out, enc_h, enc_c = encoder(inputs, enc_start_state)\n",
        "\n",
        "  dec_h = enc_h\n",
        "  dec_c = enc_c\n",
        "\n",
        "  start_tokens = tf.fill([inference_batch_size], targ_lang.word_index['<start>'])\n",
        "  end_token = targ_lang.word_index['<end>']\n",
        "\n",
        "  greedy_sampler = tfa.seq2seq.GreedyEmbeddingSampler()\n",
        "\n",
        "  # Instantiate BasicDecoder object\n",
        "  decoder_instance = tfa.seq2seq.BasicDecoder(cell=decoder.rnn_cell, sampler=greedy_sampler, output_layer=decoder.fc)\n",
        "  # Setup Memory in decoder stack\n",
        "  decoder.attention_mechanism.setup_memory(enc_out)\n",
        "\n",
        "  # set decoder_initial_state\n",
        "  decoder_initial_state = decoder.build_initial_state(inference_batch_size, [enc_h, enc_c], tf.float32)\n",
        "\n",
        "\n",
        "  ### Since the BasicDecoder wraps around Decoder's rnn cell only, you have to ensure that the inputs to BasicDecoder \n",
        "  ### decoding step is output of embedding layer. tfa.seq2seq.GreedyEmbeddingSampler() takes care of this. \n",
        "  ### You only need to get the weights of embedding layer, which can be done by decoder.embedding.variables[0] and pass this callabble to BasicDecoder's call() function\n",
        "\n",
        "  decoder_embedding_matrix = decoder.embedding.variables[0]\n",
        "  \n",
        "  outputs, _, _ = decoder_instance(decoder_embedding_matrix, start_tokens = start_tokens, end_token= end_token, initial_state=decoder_initial_state)\n",
        "  return outputs.sample_id.numpy()\n",
        "\n",
        "def translate(sentence):\n",
        "  result = evaluate_sentence(sentence)\n",
        "  print(result)\n",
        "  result = targ_lang.sequences_to_texts(result)\n",
        "  print('Input: %s' % (sentence))\n",
        "  print('Predicted translation: {}'.format(result))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n250XbnjOaqP"
      },
      "source": [
        "## Restore the latest checkpoint and test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UJpT9D5_OgP6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7f9499417390>"
            ]
          },
          "execution_count": 20,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# restoring the latest checkpoint in checkpoint_dir\n",
        "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "id": "WYmYhNN_faR5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[ 11  12  49 224  40   4   3]]\n",
            "Input: hace mucho frio aqui.\n",
            "Predicted translation: ['it s very pretty here . <end>']\n"
          ]
        }
      ],
      "source": [
        "translate(u'hace mucho frio aqui.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "id": "zSx2iM36EZQZ"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[ 20   9  22 190   4   3]]\n",
            "Input: esta es mi vida.\n",
            "Predicted translation: ['this is my life . <end>']\n"
          ]
        }
      ],
      "source": [
        "translate(u'esta es mi vida.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3LLCx3ZE0Ls"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[25  7 90  8  3]]\n",
            "Input: ¿todavia estan en casa?\n",
            "Predicted translation: ['are you home ? <end>']\n"
          ]
        }
      ],
      "source": [
        "translate(u'¿todavia estan en casa?')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DUQVLVqUE1YW"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[126  16 892  11  75   4   3]]\n",
            "Input: trata de averiguarlo.\n",
            "Predicted translation: ['try to figure it out . <end>']\n"
          ]
        }
      ],
      "source": [
        "# wrong translation\n",
        "translate(u'trata de averiguarlo.')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRUuNDeY0HiC"
      },
      "source": [
        "## Use tf-addons BeamSearchDecoder \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "id": "AJ-RTQ0hsJNL"
      },
      "outputs": [],
      "source": [
        "def beam_evaluate_sentence(sentence, beam_width=3):\n",
        "  sentence = dataset_creator.preprocess_sentence(sentence)\n",
        "\n",
        "  inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]\n",
        "  inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],\n",
        "                                                          maxlen=max_length_input,\n",
        "                                                          padding='post')\n",
        "  inputs = tf.convert_to_tensor(inputs)\n",
        "  inference_batch_size = inputs.shape[0]\n",
        "  result = ''\n",
        "\n",
        "  enc_start_state = [tf.zeros((inference_batch_size, units)), tf.zeros((inference_batch_size,units))]\n",
        "  enc_out, enc_h, enc_c = encoder(inputs, enc_start_state)\n",
        "\n",
        "  dec_h = enc_h\n",
        "  dec_c = enc_c\n",
        "\n",
        "  start_tokens = tf.fill([inference_batch_size], targ_lang.word_index['<start>'])\n",
        "  end_token = targ_lang.word_index['<end>']\n",
        "\n",
        "  # From official documentation\n",
        "  # NOTE If you are using the BeamSearchDecoder with a cell wrapped in AttentionWrapper, then you must ensure that:\n",
        "  # The encoder output has been tiled to beam_width via tfa.seq2seq.tile_batch (NOT tf.tile).\n",
        "  # The batch_size argument passed to the get_initial_state method of this wrapper is equal to true_batch_size * beam_width.\n",
        "  # The initial state created with get_initial_state above contains a cell_state value containing properly tiled final state from the encoder.\n",
        "\n",
        "  enc_out = tfa.seq2seq.tile_batch(enc_out, multiplier=beam_width)\n",
        "  decoder.attention_mechanism.setup_memory(enc_out)\n",
        "  print(\"beam_with * [batch_size, max_length_input, rnn_units] :  3 * [1, 16, 1024]] :\", enc_out.shape)\n",
        "\n",
        "  # set decoder_inital_state which is an AttentionWrapperState considering beam_width\n",
        "  hidden_state = tfa.seq2seq.tile_batch([enc_h, enc_c], multiplier=beam_width)\n",
        "  decoder_initial_state = decoder.rnn_cell.get_initial_state(batch_size=beam_width*inference_batch_size, dtype=tf.float32)\n",
        "  decoder_initial_state = decoder_initial_state.clone(cell_state=hidden_state)\n",
        "\n",
        "  # Instantiate BeamSearchDecoder\n",
        "  decoder_instance = tfa.seq2seq.BeamSearchDecoder(decoder.rnn_cell,beam_width=beam_width, output_layer=decoder.fc)\n",
        "  decoder_embedding_matrix = decoder.embedding.variables[0]\n",
        "\n",
        "  # The BeamSearchDecoder object's call() function takes care of everything.\n",
        "  outputs, final_state, sequence_lengths = decoder_instance(decoder_embedding_matrix, start_tokens=start_tokens, end_token=end_token, initial_state=decoder_initial_state)\n",
        "  # outputs is tfa.seq2seq.FinalBeamSearchDecoderOutput object. \n",
        "  # The final beam predictions are stored in outputs.predicted_id\n",
        "  # outputs.beam_search_decoder_output is a tfa.seq2seq.BeamSearchDecoderOutput object which keep tracks of beam_scores and parent_ids while performing a beam decoding step\n",
        "  # final_state = tfa.seq2seq.BeamSearchDecoderState object.\n",
        "  # Sequence Length = [inference_batch_size, beam_width] details the maximum length of the beams that are generated\n",
        "\n",
        "  \n",
        "  # outputs.predicted_id.shape = (inference_batch_size, time_step_outputs, beam_width)\n",
        "  # outputs.beam_search_decoder_output.scores.shape = (inference_batch_size, time_step_outputs, beam_width)\n",
        "  # Convert the shape of outputs and beam_scores to (inference_batch_size, beam_width, time_step_outputs)\n",
        "  final_outputs = tf.transpose(outputs.predicted_ids, perm=(0,2,1))\n",
        "  beam_scores = tf.transpose(outputs.beam_search_decoder_output.scores, perm=(0,2,1))\n",
        "  \n",
        "  return final_outputs.numpy(), beam_scores.numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "id": "g_LvXGvX8X-O"
      },
      "outputs": [],
      "source": [
        "def beam_translate(sentence):\n",
        "  result, beam_scores = beam_evaluate_sentence(sentence)\n",
        "  print(result.shape, beam_scores.shape)\n",
        "  for beam, score in zip(result, beam_scores):\n",
        "    print(beam.shape, score.shape)\n",
        "    output = targ_lang.sequences_to_texts(beam)\n",
        "    output = [a[:a.index('<end>')] for a in output]\n",
        "    beam_score = [a.sum() for a in score]\n",
        "    print('Input: %s' % (sentence))\n",
        "    for i in range(len(output)):\n",
        "      print('{} Predicted translation: {}  {}'.format(i+1, output[i], beam_score[i]))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "id": "TODnXBleDzzO"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "beam_with * [batch_size, max_length_input, rnn_units] :  3 * [1, 16, 1024]] : (3, 16, 1024)\n",
            "(1, 3, 7) (1, 3, 7)\n",
            "(3, 7) (3, 7)\n",
            "Input: hace mucho frio aqui.\n",
            "1 Predicted translation: it s very pretty here .   -4.117094039916992\n",
            "2 Predicted translation: it s very cold here .   -14.85302734375\n",
            "3 Predicted translation: it s very pretty news .   -25.59416389465332\n"
          ]
        }
      ],
      "source": [
        "beam_translate(u'hace mucho frio aqui.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_BezQwENFY3L"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "beam_with * [batch_size, max_length_input, rnn_units] :  3 * [1, 16, 1024]] : (3, 16, 1024)\n",
            "(1, 3, 7) (1, 3, 7)\n",
            "(3, 7) (3, 7)\n",
            "Input: ¿todavia estan en casa?\n",
            "1 Predicted translation: are you still home ?   -4.036754131317139\n",
            "2 Predicted translation: are you still at home ?   -15.306867599487305\n",
            "3 Predicted translation: are you still go home ?   -20.533388137817383\n"
          ]
        }
      ],
      "source": [
        "beam_translate(u'¿todavia estan en casa?')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "networks_seq2seq_nmt.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
